#
# example032.py
#
# The wavelet transform: Q15 format
#
# Copyright (C) 2012 Robert Buj Gelonch
# Copyright (C) 2012 David Megias Jimenez
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
__author__ = "Robert Buj Gelonch, and David Megias Jimenez"
__copyright__ = "Copyright 2012, Robert Buj Gelonch and David Megias Jimenez"
__credits__ = ["Robert Buj Gelonch", "David Megias Jimenez"]
__license__ = "GPL"
__version__ = "3"
__maintainer__ = "Robert Buj"
__email__ = "rbuj@uoc.edu"
__status__ = "Development"
__docformat__ = 'plaintext'

from numpy import arange
from numpy import sinc
from numpy import zeros
from pylab import figure
from pylab import plot
from pylab import show
from pylab import stem
from pylab import subplot
from pylab import title
from pywt import Wavelet
from pywt import dwt
from scipy.signal import convolve
from numpy import ceil

print "example032.py"
print
print "The discrete wavelet transform"
print
#--------------------------------------------------------------
# Signsl
#--------------------------------------------------------------
x = sinc(arange(-20., 21.) / 5)
#--------------------------------------------------------------
# Wavelet Transform
#--------------------------------------------------------------
# zpd (zero-padding): signal is extended by adding zero samples
# cpd (constant-padding): border values are replicated
# sym (symmetric-padding): signal is extended by mirroring samples
# ppd (periodic-padding): signal is treated as a periodic one
# sp1 (smooth-padding): signal is extended according to the first
#      derivatives calculated on the edges (straight line)
# per (periodization): is like periodic-padding but gives the
#     smallest possible number of decomposition coefficients.
#     IDWT must be performed with the same mode
# http://www.pybytes.com/pywavelets/ref/signal-extension-modes.html#modes
# http://wavelets.pybytes.com/
#--------------------------------------------------------------
mode = 'zpd'
wavelet1 = Wavelet('db10')
cA, cD = dwt(x, wavelet1, mode)

################################################
# Filter coefficients / impulse response
################################################
# http://www.pybytes.com/pywavelets/ref/wavelets.html

dec_lo_Q15 = [ceil(coeff/(2**(-15)))*(2**(-15)) for coeff in wavelet1.dec_lo]
dec_hi_Q15 = [ceil(coeff/(2**(-15)))*(2**(-15)) for coeff in wavelet1.dec_hi]
rec_lo_Q15 = [ceil(coeff/(2**(-15)))*(2**(-15)) for coeff in wavelet1.rec_lo]
rec_hi_Q15 = [ceil(coeff/(2**(-15)))*(2**(-15)) for coeff in wavelet1.rec_hi]

fig = figure(num=None, \
             figsize=(14, 9), \
             dpi=80, \
             facecolor='w', \
             edgecolor='k')
fig.suptitle("Filter coefficients: Q15 format")
ax2 = subplot(221)
title(r'Decomposition LP filter')
stem(range(len(dec_lo_Q15)), dec_lo_Q15)
ax2 = subplot(222)
title(r'Decomposition HP filter')
stem(range(len(dec_hi_Q15)), dec_hi_Q15)
ax2 = subplot(223)
title(r'Reconstruction LP filter')
stem(range(len(rec_lo_Q15)), rec_lo_Q15)
ax2 = subplot(224)
title(r'Reconstruction HP filter')
stem(range(len(rec_hi_Q15)), rec_hi_Q15)

################################################
# time-domain filtering
################################################
x_Q15 = [ceil(elem/(2**(-15)))*(2**(-15)) for elem in x]
convolve_dec_low = convolve(x_Q15, dec_lo_Q15, mode='same')
convolve_dec_hight = convolve(x_Q15, dec_hi_Q15, mode='same')
cA = [convolve_dec_low[2 * i] for i in range(len(convolve_dec_low) / 2)]
cA = [ceil(elem/(2**(-15)))*(2**(-15)) for elem in cA]
cD = [convolve_dec_hight[2 * i] for i in range(len(convolve_dec_hight) / 2)]
cD = [ceil(elem/(2**(-15)))*(2**(-15)) for elem in cD]
A = zeros(len(cA) * 2)
D = zeros(len(cD) * 2)
for i in range(len(cA)):
    A[2 * i] = cA[i]
for i in range(len(cD)):
    D[2 * i] = cD[i]
s = convolve(A, rec_lo_Q15, mode='same') + convolve(D, rec_hi_Q15, mode='same')
s = ceil(s/(2**(-15)))*(2**(-15))

fig = figure(num=2, figsize=(14, 9), dpi=80, facecolor='w', edgecolor='k')
fig.suptitle("DWT Time Domain: Q15 format")

ax2 = subplot(321)
title("Original signal")
plot(x, 'go--', label="x")
leg = ax2.legend(loc='best', fancybox=True)
leg.get_frame().set_alpha(0.5)

ax2 = subplot(323)
title("Descomposition: convolution")
plot(convolve_dec_low, 'ro--', label="(x*g)")
plot(convolve_dec_hight, 'bo--', label="(x*h)")
leg = ax2.legend(loc='best', fancybox=True)
leg.get_frame().set_alpha(0.5)

ax2 = subplot(325)
title("Descomposition: Downsample")
plot(cA, 'ro--', label="(x*g)" + unichr(8595) + "2")
plot(cD, 'bo--', label="(x*h)" + unichr(8595) + "2")
leg = ax2.legend(loc='best', fancybox=True)
leg.get_frame().set_alpha(0.5)

ax2 = subplot(322)
title("Recomposition: upsample & convolution")
plot(s, 'ro--', label="(x*g)" + unichr(8593) + "2")
leg = ax2.legend(loc='best', fancybox=True)
leg.get_frame().set_alpha(0.5)


show()

print "Done"
